#include "marching_squares.h"
#include "data_structures.hpp"
#include <cmath>
#include <algorithm>
#include <iostream>
#include <set>
#include <fstream>
#include <filesystem>
#include <chrono>
#include <iomanip>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

namespace np {

    Eigen::VectorXd linspace(double start, double end, int num) {
        Eigen::VectorXd vec(num);
        if (num == 1) {
            vec[0] = start;
            return vec;
        }
        double step = (end - start) / (num - 1);
        for (int i = 0; i < num; ++i) {
            vec[i] = start + i * step;
        }
        return vec;
    }

    Eigen::VectorXd arange(int size) {
        Eigen::VectorXd vec(size);
        for (int i = 0; i < size; ++i) {
            vec[i] = i;
        }
        return vec;
    }

    double interp(double x, const Eigen::VectorXd& xp, const Eigen::VectorXd& fp) {
        if (x <= xp[0]) return fp[0];
        if (x >= xp[xp.size() - 1]) return fp[fp.size() - 1];

        int i = 0;
        while (x > xp[i+1]) i++;

        double x1 = xp[i];
        double y1 = fp[i];
        double x2 = xp[i+1];
        double y2 = fp[i+1];

        return y1 + (y2 - y1) * (x - x1) / (x2 - x1);
    }
}

namespace EddyDetection {

namespace {

    double calculatePercentileThreshold(const Eigen::MatrixXd& field, double percentile,
                                        bool positive_values_only = false) {
        std::vector<double> values;
        values.reserve(field.rows() * field.cols());
        
        for (int i = 0; i < field.rows(); ++i) {
            for (int j = 0; j < field.cols(); ++j) {
                if (std::isfinite(field(i, j))) {
                    if (positive_values_only && field(i, j) <= 0) continue;
                    if (!positive_values_only && field(i, j) >= 0) continue;
                    values.push_back(field(i, j));
                }
            }
        }
        
        if (values.empty()) return 0.0;
        
        std::sort(values.begin(), values.end());
        size_t index = static_cast<size_t>(percentile / 100.0 * (values.size() - 1));

        return values[index];
    }

    void floodFill(const Eigen::MatrixXi& binary_image, Eigen::MatrixXi& labels,
                   int x, int y, int label, std::vector<std::pair<int, int>>& region_pixels) {
        if (x < 0 || x >= binary_image.rows() || y < 0 || y >= binary_image.cols() ||
            labels(x, y) != 0 || binary_image(x, y) == 0) {
            return;
        }
        
        labels(x, y) = label;
        region_pixels.push_back({x, y});
        
        floodFill(binary_image, labels, x+1, y, label, region_pixels);
        floodFill(binary_image, labels, x-1, y, label, region_pixels);
        floodFill(binary_image, labels, x, y+1, label, region_pixels);
        floodFill(binary_image, labels, x, y-1, label, region_pixels);
        floodFill(binary_image, labels, x+1, y+1, label, region_pixels);
        floodFill(binary_image, labels, x-1, y-1, label, region_pixels);
        floodFill(binary_image, labels, x+1, y-1, label, region_pixels);
        floodFill(binary_image, labels, x-1, y+1, label, region_pixels);
    }
    
    std::vector<Eigen::Vector2d> extractBoundary(const std::vector<std::pair<int, int>>& region_pixels,
                                                  const Eigen::VectorXd& lat, const Eigen::VectorXd& lon) {
        std::vector<Eigen::Vector2d> boundary;
        
        if (region_pixels.empty()) return boundary;
        
        int min_x = region_pixels[0].first, max_x = region_pixels[0].first;
        int min_y = region_pixels[0].second, max_y = region_pixels[0].second;
        
        for (const auto& pixel : region_pixels) {
            min_x = std::min(min_x, pixel.first);
            max_x = std::max(max_x, pixel.first);
            min_y = std::min(min_y, pixel.second);
            max_y = std::max(max_y, pixel.second);
        }
        
        int width = max_x - min_x + 3;
        int height = max_y - min_y + 3;
        
        Eigen::MatrixXi mask = Eigen::MatrixXi::Zero(width, height);
        
        for (const auto& pixel : region_pixels) {
            int local_x = pixel.first - min_x + 1;
            int local_y = pixel.second - min_y + 1;
            if (local_x >= 0 && local_x < width && local_y >= 0 && local_y < height) {
                mask(local_x, local_y) = 1;
            }
        }
        
        for (const auto& pixel : region_pixels) {
            int x = pixel.first;
            int y = pixel.second;
            int local_x = x - min_x + 1;
            int local_y = y - min_y + 1;
            
            bool is_boundary = false;
            
            for (int dx = -1; dx <= 1; ++dx) {
                for (int dy = -1; dy <= 1; ++dy) {
                    int nx = local_x + dx;
                    int ny = local_y + dy;
                    if (nx >= 0 && nx < width && ny >= 0 && ny < height) {
                        if (mask(nx, ny) == 0) {
                            is_boundary = true;
                            break;
                        }
                    } else {
                        is_boundary = true;
                        break;
                    }
                }
                if (is_boundary) break;
            }
            
            if (is_boundary) {
                if (x >= 0 && x < lat.size() && y >= 0 && y < lon.size()) {
                    boundary.push_back(Eigen::Vector2d(lat[x], lon[y]));
                }
            }
        }
        
        return boundary;
    }

    // Vorticity calculation, considering spherical coordinates
    Eigen::MatrixXd compute_vorticity_spherical(const Eigen::VectorXd& lat, const Eigen::VectorXd& lon, 
                                                const Eigen::MatrixXd& u, const Eigen::MatrixXd& v) {
        if (u.rows() != lat.size() || u.cols() != lon.size()) {
             throw std::runtime_error("The velocity field dimension does not match the coordinate dimension.");
        }
        
        const double R_EARTH = 6371e3;
        
        Eigen::MatrixXd vorticity(lat.size(), lon.size());

        Eigen::MatrixXd dv_dx = Eigen::MatrixXd::Zero(lat.size(), lon.size());
        Eigen::MatrixXd du_dy = Eigen::MatrixXd::Zero(lat.size(), lon.size());

        for(int i = 0; i < lat.size(); ++i) {
            double cos_lat = std::cos(lat(i) * M_PI / 180.0);
            if (cos_lat < 1e-9) cos_lat = 1e-9;
            double circle_lat = 2.0 * M_PI * R_EARTH * cos_lat;
            
            for(int j = 1; j < lon.size() - 1; ++j) {
                dv_dx(i,j) = (v(i, j+1) - v(i, j-1)) / (circle_lat * (lon(j+1) - lon(j-1)) / 360.0);
            }
        }

        double circle_lon = 2.0 * M_PI * R_EARTH;
        for(int j = 0; j < lon.size(); ++j) {
            for(int i = 1; i < lat.size() - 1; ++i) {
                du_dy(i,j) = (u(i+1, j) - u(i-1, j)) / (circle_lon * (lat(i+1) - lat(i-1)) / 360.0);
            }
        }

        vorticity = dv_dx - du_dy;
        return vorticity;
    }

    double calculate_area_km2(const std::vector<Eigen::Vector2d>& boundary_coords, const Eigen::Vector2d& center) {
        if (boundary_coords.size() < 3) return 0.0;
        // center[0] is latitude, center[1] is longitude
        double mean_lat_rad = center[0] * M_PI / 180.0;
        double km_per_deg_lon = 111.32 * std::cos(mean_lat_rad);
        double km_per_deg_lat = 111.32;
        double area = 0.0;
        for (size_t i = 0; i < boundary_coords.size(); ++i) {
            int j = (i + 1) % boundary_coords.size();
            // boundary_coords[i][0] is latitude, boundary_coords[i][1] is longitude
            double x1 = (boundary_coords[i][1] - center[1]) * km_per_deg_lon;
            double y1 = (boundary_coords[i][0] - center[0]) * km_per_deg_lat;
            double x2 = (boundary_coords[j][1] - center[1]) * km_per_deg_lon;
            double y2 = (boundary_coords[j][0] - center[0]) * km_per_deg_lat;
            area += x1 * y2 - x2 * y1;
        }
        return std::abs(area) / 2.0;
    }

    std::vector<std::vector<std::pair<int, int>>> find_connected_components(const Eigen::MatrixXi& binary_map) {
        Eigen::MatrixXi visited = Eigen::MatrixXi::Zero(binary_map.rows(), binary_map.cols());
        std::vector<std::vector<std::pair<int, int>>> components;

        for (int i = 0; i < binary_map.rows(); ++i) {
            for (int j = 0; j < binary_map.cols(); ++j) {
                if (binary_map(i, j) == 1 && visited(i, j) == 0) {
                    std::vector<std::pair<int, int>> current_component;
                    std::vector<std::pair<int, int>> queue;
                    
                    queue.push_back({i, j});
                    visited(i, j) = 1;

                    int head = 0;
                    while(head < queue.size()) {
                        std::pair<int, int> p = queue[head++];
                        current_component.push_back(p);
                        
                        for (int dr = -1; dr <= 1; ++dr) {
                            for (int dc = -1; dc <= 1; ++dc) {
                                if (dr == 0 && dc == 0) continue;
                                int nr = p.first + dr;
                                int nc = p.second + dc;

                                if (nr >= 0 && nr < binary_map.rows() && nc >= 0 && nc < binary_map.cols() &&
                                    binary_map(nr, nc) == 1 && visited(nr, nc) == 0) {
                                    visited(nr, nc) = 1;
                                    queue.push_back({nr, nc});
                                }
                            }
                        }
                    }
                    components.push_back(current_component);
                }
            }
        }
        return components;
    }

    std::vector<Eigen::Vector2d> extract_boundary(
        const std::vector<std::pair<int, int>>& component,
        const Eigen::VectorXd& lat, const Eigen::VectorXd& lon) {
        
        if (component.size() < 4) return {};
        
        std::set<std::pair<int, int>> component_set(component.begin(), component.end());
        
        std::vector<Eigen::Vector2d> boundary_coords;
        
        for (const auto& p : component) {
            bool is_boundary = false;
            int dr[] = {0, 0, 1, -1};
            int dc[] = {1, -1, 0, 0};
            
            for (int i = 0; i < 4; ++i) {
                int nr = p.first + dr[i];
                int nc = p.second + dc[i];
                if (component_set.find({nr, nc}) == component_set.end()) {
                    is_boundary = true;
                    break;
                }
            }
            
            if (is_boundary) {
                boundary_coords.push_back({ lat(p.first), lon(p.second)});
            }
        }
        
        if (boundary_coords.size() > 2) {
            Eigen::Vector2d center = {0, 0};
            for(const auto& p : boundary_coords) center += p;
            center /= boundary_coords.size();
            
            std::sort(boundary_coords.begin(), boundary_coords.end(), 
                [&](const Eigen::Vector2d& a, const Eigen::Vector2d& b) {
                return std::atan2(a.x() - center.x(), a.y() - center.y()) <
                       std::atan2(b.x() - center.x(), b.y() - center.y());
            });
        }

        return boundary_coords;
    }

    std::vector<EddyContour> find_and_analyze_regions(
        const Eigen::MatrixXi& binary_map, const Eigen::MatrixXd& vorticity, 
        const Eigen::MatrixXd& temp_data, const Eigen::VectorXd& lat, const Eigen::VectorXd& lon, 
        const EddyDetectionConfig& config, const EddyContourDataCyclone& eddy_type) {
        
        std::vector<EddyContour> eddies;
        auto components = find_connected_components(binary_map);

        for (const auto& component : components) {
            if (component.empty()) continue;

            double sum_lat_idx = 0, sum_lon_idx = 0, sum_vort = 0;
            for (const auto& p : component) {
                sum_lat_idx += p.first;
                sum_lon_idx += p.second;
                sum_vort += vorticity(p.first, p.second);
            }

            Eigen::Vector2d center_idx(sum_lat_idx / component.size(), sum_lon_idx / component.size());
            Eigen::Vector2d center_coord(
                np::interp(center_idx[0], np::arange(lat.size()), lat),
                np::interp(center_idx[1], np::arange(lon.size()), lon)
            );

            auto boundary_coords = extract_boundary(component, lat, lon);
            if (boundary_coords.empty()) continue;
            
            double area_km2 = calculate_area_km2(boundary_coords, center_coord);
            
            if (area_km2 >= config.min_area_km2) {
                EddyContour eddy;
                eddy.center = center_coord;
                eddy.mean_vorticity = sum_vort / component.size();
                eddy.mean_omega = fabs(eddy.mean_vorticity/2);
                eddy.type = eddy_type;
                eddy.boundary_points = boundary_coords;
                eddy.area_km2 = area_km2;
                
                double temp_inside = 0;
                for(const auto& p : component) temp_inside += temp_data(p.first, p.second);
                temp_inside /= component.size();

                double temp_boundary = 0;
                for(const auto& p : boundary_coords) temp_boundary += 
                    np::interp(
                        p.y(), lon, temp_data.row(
                            np::interp(
                                p.x(), lat, np::arange(lat.size())
                            )
                        )
                    );
                if(!boundary_coords.empty()) temp_boundary /= boundary_coords.size();

                eddy.temp_anomaly = temp_inside - temp_boundary;
                if (eddy.temp_anomaly > 0) {
                    eddy.thermal_type = EddyContourDataThermal::WarmCore;
                } else {
                    eddy.thermal_type = EddyContourDataThermal::Cold_core;
                }
                
                eddies.push_back(eddy);
            }
        }
        return eddies;
    }
}

Eigen::MatrixXd EulerianCore::computeVorticity(const VelocityField& velocity_field, int time_index) {

    if (velocity_field.u_layers.empty() || velocity_field.v_layers.empty()) {
        std::cerr << "Error: Speed field data is empty" << std::endl;
        return Eigen::MatrixXd();
    }
    
    const auto& u = velocity_field.u_layers[time_index < velocity_field.u_layers.size() ? time_index : 0];
    const auto& v = velocity_field.v_layers[time_index < velocity_field.v_layers.size() ? time_index : 0];
    
    return compute_vorticity_spherical(velocity_field.lat, velocity_field.lon, u, v);
}

std::vector<std::vector<Eigen::Vector2d>> EulerianCore::findVorticityContours(
    const Eigen::MatrixXd& vorticity_field,
    const Eigen::VectorXd& lat, const Eigen::VectorXd& lon,
    double threshold_percentile, double min_area_km2) {
    
    std::vector<std::vector<Eigen::Vector2d>> contours;
    
    
    // Calculate the area of grid cells
    double mean_lat_rad = lat.mean() * M_PI / 180.0;
    double dx_km = (lon.size() > 1 ? (lon[1] - lon[0]) : 1.0) * 111.32 * std::cos(mean_lat_rad);
    double dy_km = (lat.size() > 1 ? (lat[1] - lat[0]) : 1.0) * 111.32;
    double pixel_area_km2 = dx_km * dy_km;
    int min_pixels = static_cast<int>(min_area_km2 / pixel_area_km2);
    
    // Dealing with cyclonic eddies (positive vorticity)
    double cyclo_threshold = calculatePercentileThreshold(vorticity_field, threshold_percentile, true);
    
    Eigen::MatrixXi cyclo_binary = Eigen::MatrixXi::Zero(vorticity_field.rows(), vorticity_field.cols());
    for (int i = 0; i < vorticity_field.rows(); ++i) {
        for (int j = 0; j < vorticity_field.cols(); ++j) {
            cyclo_binary(i, j) = (vorticity_field(i, j) >= cyclo_threshold) ? 1 : 0;
        }
    }

    // Mark connected regions
    Eigen::MatrixXi labels = Eigen::MatrixXi::Zero(cyclo_binary.rows(), cyclo_binary.cols());
    int current_label = 1;
    
    for (int i = 0; i < cyclo_binary.rows(); ++i) {
        for (int j = 0; j < cyclo_binary.cols(); ++j) {
            if (cyclo_binary(i, j) == 1 && labels(i, j) == 0) {
                std::vector<std::pair<int, int>> region_pixels;
                floodFill(cyclo_binary, labels, i, j, current_label, region_pixels);
                
                if (region_pixels.size() >= min_pixels) {
                    auto boundary = extractBoundary(region_pixels, lat, lon);    

                    if (boundary.size() >= 4) {
                        contours.push_back(boundary);
                    }
                }
                current_label++;
            }
        }
    }

    // Dealing with anticyclonic eddies (negative vorticity)
    double anti_threshold = calculatePercentileThreshold(vorticity_field, 100.0 - threshold_percentile, false);
    
    Eigen::MatrixXi anti_binary = Eigen::MatrixXi::Zero(vorticity_field.rows(), vorticity_field.cols());
    for (int i = 0; i < vorticity_field.rows(); ++i) {
        for (int j = 0; j < vorticity_field.cols(); ++j) {
            anti_binary(i, j) = (vorticity_field(i, j) <= anti_threshold) ? 1 : 0;
        }
    }
    
    labels = Eigen::MatrixXi::Zero(anti_binary.rows(), anti_binary.cols());
    
    for (int i = 0; i < anti_binary.rows(); ++i) {
        for (int j = 0; j < anti_binary.cols(); ++j) {
            if (anti_binary(i, j) == 1 && labels(i, j) == 0) {
                std::vector<std::pair<int, int>> region_pixels;
                floodFill(anti_binary, labels, i, j, current_label, region_pixels);
                
                if (region_pixels.size() >= min_pixels) {
                    auto boundary = extractBoundary(region_pixels, lat, lon);
                    if (boundary.size() >= 4) {
                        contours.push_back(boundary);
                    }
                }
                current_label++;
            }
        }
    }
    
    return contours;
}

std::vector<std::vector<Eigen::Vector2d>> EulerianCore::findVorticityContours1(
    const Eigen::MatrixXd& vorticity_field,
    const Eigen::VectorXd& lat, const Eigen::VectorXd& lon,
    double threshold_percentile, double min_area_km2) {

    // Dealing with cyclonic eddies (positive vorticity)
    double cyclo_threshold = calculatePercentileThreshold(vorticity_field, threshold_percentile, true);

    Eigen::MatrixXi cyclo_binary = Eigen::MatrixXi::Zero(vorticity_field.rows(), vorticity_field.cols());
    for (int i = 0; i < vorticity_field.rows(); ++i) {
        for (int j = 0; j < vorticity_field.cols(); ++j) {
            cyclo_binary(i, j) = (vorticity_field(i, j) >= cyclo_threshold) ? 1 : 0;
        }
    }

    // Dealing with anticyclonic eddies (negative vorticity)
    double anti_threshold = calculatePercentileThreshold(vorticity_field, 100.0 - threshold_percentile, false);

    Eigen::MatrixXi anti_binary = Eigen::MatrixXi::Zero(vorticity_field.rows(), vorticity_field.cols());
    for (int i = 0; i < vorticity_field.rows(); ++i) {
        for (int j = 0; j < vorticity_field.cols(); ++j) {
            anti_binary(i, j) = (vorticity_field(i, j) <= anti_threshold) ? 1 : 0;
        }
    }

    std::vector<std::vector<Eigen::Vector2d>> cycloContours, antiContours, contours;
    MarchingSquares ms;
    cycloContours = ms.extractContours(cyclo_binary, lat, lon);
    antiContours = ms.extractContours(anti_binary, lat, lon);

    contours.insert(contours.end(), cycloContours.begin(), cycloContours.end());
    contours.insert(contours.end(), antiContours.begin(), antiContours.end());

    return contours;
}


std::vector<EddyContour> EulerianCore::detectEddiesEulerian(
    const VelocityField& velocity_field,
    const TemperatureField& temp_field,
    const EddyDetectionConfig& config,
    int time_index) {
    
    //std::cout << "Start Euler Eddy detection..." << std::endl;

    if (velocity_field.u_layers.empty() || velocity_field.v_layers.empty()) {
        std::cerr << "Error: Speed field data is empty" << std::endl;
        return std::vector<EddyContour>();
    }

    Eigen::MatrixXd vorticity = computeVorticity(velocity_field, time_index);  // calculate vorticity

    auto contours = findVorticityContours1(vorticity, velocity_field.lat, velocity_field.lon,  ////////
                                          config.vorticity_threshold_percentile, config.min_area_km2);
    
    //std::cout << "Find " << contours.size() << " Vorticity profile" << std::endl;

    std::vector<EddyContour> eddies;
    int i = 0;
    for (const auto& contour : contours) {

        EddyContour eddy = EddyAnalyzer::analyzeEddy(contour, vorticity, temp_field.data_layers[0],
                                                     velocity_field.lat, velocity_field.lon,
                                                     config.min_area_km2);
        if (eddy.area_km2 >= config.min_area_km2) {
            eddies.push_back(eddy);
        }
    }

    eddies = EddyAnalyzer::filterAndSortEddies(eddies, config.min_area_km2, velocity_field);
    
    //std::cout << "Have Detect " << eddies.size() << " Effective Eddy" << std::endl;
    return eddies;
}

} // namespace EddyDetection 